import argparse
import sys
import os
import gzip
import mmap
import numpy as np
import time
from multiprocessing import Manager
import multiprocessing as mp

def generate_key(line):
    # 根据前两列生成key
    columns = line.split()[:2]
    return '||'.join(columns)

def process_file(output_file):
     # 检查输出文件是否存在
    try:
        with open(output_file, 'r') as f:
            existing_keys = set()
            for line in f:
                existing_keys.add(generate_key(line))
    except FileNotFoundError:
        existing_keys = set()
    return existing_keys

def replace_and_clean(string):
    # 将'], '替换成';'
    string = string.replace('],[', ';')
    # 删除空格和左方括号
    string = string.replace(' ', '').replace('[', '').replace(']', '')#.replace('\'','')
    return string

def read_feature(features_file,output_Q,existing_keys,start_pos,end_pos, batch_size=10):
    if features_file.endswith(".gz"):
        infile = gzip.open(features_file, 'rt')
    else:
        infile = open(features_file, 'r')
    fea_list=[]
    with mmap.mmap(infile.fileno(), 0, access=mmap.ACCESS_READ) as mm:
        # 在字节块的开始位置搜索换行符
        mm.seek(start_pos)
        start = mm.tell()
        if start != 0 and mm[start-1] != b'\n':
            start = mm.find(b'\n', 0, start)
            if start != -1:
                start += 1  # 跳过换行符
            else:
                #start_pos = 0
                
                sys.exit()
            mm.seek(start)
        
        ## 如果起始位置不是行的开头，则向后移动，直到找到行的开头
        #while mm.tell() < end_pos and mm.read_byte() != b'\n':
        #    mm.seek(1, 1) 
        #start = mm.tell()
        while mm.tell() < end_pos:
            line = mm.readline()
            if not line:
                break
            try:
                decoded_line = line.decode('cp1252')  # 尝试使用 utf-8 编码解码
            except UnicodeDecodeError:
                decoded_line = line.decode('latin-1')
            words=decoded_line.strip().split('\t')            
            if len(words)!=11:
                print('words is less than 11 with {}'.format(words[0]))
                continue
            if int(words[10])<10:
                print('mapq is too low with {}'.format(words[0]))
                continue
            key='||'.join(words[:2])
            if key in existing_keys:
                continue
            #print(words)
            #break
            norm_signals=replace_and_clean(words[4]).split(";")
            #print(norm_signals[0])
            #print(norm_signals[1])
            norm_signals_new=[]
            for sig in norm_signals:
                norm_signals_new.append(np.round([float(y) for y in sig.split(",")], decimals=6))
            try:
                pred_dorado=np.round(np.array([float(x) for x in words[6].split(",")]), decimals=6)
            except:
                print(words[6])
                continue
            try:
                pred_deepsignal=np.round(np.array([float(x) for x in words[7].split(",")]), decimals=6)
            except:
                print(words[7])
                continue
            try:
                pred_mean=np.round(np.array([float(x) for x in words[8].split(",")]), decimals=6)
            except:
                print(words[8])
                continue
            pred_dorado_text=','.join([str(x) for x in pred_dorado])
            pred_deepsignal_text=','.join([str(x) for x in pred_deepsignal])
            pred_mean_text=','.join([str(x) for x in pred_mean])
            norm_signals_text=';'.join([",".join([str(y) for y in x]) for x in norm_signals_new])
            fea_str='\t'.join([words[0],words[1],words[2],words[3],norm_signals_text,words[5],pred_dorado_text,pred_deepsignal_text,pred_mean_text,words[9],words[10]])
            fea_list.append(fea_str)
            if len(fea_list)>=batch_size:
                output_Q.put(fea_list)
                fea_list=[]
        if len(fea_list)>0:
            output_Q.put(fea_list)


def _write_featurestr(write_fp, featurestr_q,time_wait = 3):
    print('write_process-{} starts'.format(os.getpid()))
    output=0
    if os.path.exists(write_fp):
        with open(write_fp, 'a') as wf:
            while True:
                # during test, it's ok without the sleep(time_wait)
                if featurestr_q.empty():
                    time.sleep(time_wait)
                    continue
                features_str = featurestr_q.get()
                if features_str == "kill":
                    print('output line {}'.format(output))
                    print('write_process-{} finished'.format(os.getpid()))
                    break
                for one_features_str in features_str:
                    output+=1
                    wf.write(one_features_str + "\n")
                wf.flush()
    else:
        with open(write_fp, 'w') as wf:
            while True:
                # during test, it's ok without the sleep(time_wait)
                if featurestr_q.empty():
                    time.sleep(time_wait)
                    continue
                features_str = featurestr_q.get()
                if features_str == "kill":
                    print('output line {}'.format(output))
                    print('write_process-{} finished'.format(os.getpid()))
                    break
                for one_features_str in features_str:
                    output+=1
                    wf.write(one_features_str + "\n")
                wf.flush()

    

def parse_args():
    parser = argparse.ArgumentParser("")
    parser.add_argument("--input", type=str,required=True)
    parser.add_argument("--output", type=str,required=True)
    parser.add_argument("--nproc", "-p", type=int,required=True)
    parser.add_argument("--d_batch_size", action="store", type=int, default=10,
                         required=False)

    return parser.parse_args()


def main():
    args=parse_args()
    #extract(args)
    manager = mp.Manager()
    output_Q = manager.Queue()
    rnproc = args.nproc-1
    r_procs = []
    existing_keys=process_file(args.output)
    total_size=os.path.getsize(args.input)
    ex_size=int(total_size/rnproc)
    for i in range(rnproc):
        start_pos = i * ex_size
        end_pos = (i + 1) * ex_size if i < rnproc - 1 else total_size
        p_rf = mp.Process(target=read_feature, args=(args.input,output_Q,existing_keys,start_pos,end_pos,
                                                            args.d_batch_size),
                          name="reader")
            
        p_rf.daemon = True
        p_rf.start()
        r_procs.append(p_rf)
    p_w = mp.Process(target=_write_featurestr, args=(args.output, output_Q),
                         name="writer")
    p_w.daemon = True
    p_w.start()
    for p_rf in r_procs:
            p_rf.join()
    output_Q.put("kill")
    p_w.join()

if __name__ == '__main__':
    sys.exit(main())